from oasr.recognize.ctc import CTCRecognizer
from oasr.recognize.speech2text import SpeechToTextRecognizer
from oasr.recognize.transducer import TransducerRecognizer
from oasr.recognize.fastskip import FastSkipTransducerRecognizer


def build_recognizer(model_type, model, lm, args, idx2unit, output_dir):
    if model_type == 'speech2text':
        return SpeechToTextRecognizer(
            model=model, lm=lm, lm_weight=args.lm_weight,
            ctc_weight=args.ctc_weight, beam_width=args.beam_width, nbest=args.nbest, max_len=args.max_len,
            idx2unit=idx2unit, penalty=args.penalty, lamda=args.lamda, ngpu=args.ngpu)
    elif model_type == 'ctc':
        return CTCRecognizer(
            model=model, lm=lm, lm_weight=args.lm_weight, ngram_lm=args.ngram_lm, beam_width=args.beam_width,
            idx2unit=idx2unit, ngpu=args.ngpu, mode=args.mode, alpha=args.alpha, beta=args.beta,
            is_visualized=args.is_visualized, output_dir=output_dir)
    elif model_type == 'transducer':
        return TransducerRecognizer(
            model=model, lm=lm, lm_weight=args.lm_weight, beam_width=args.beam_width, max_len=args.max_len, nbest=args.nbest,
            idx2unit=idx2unit, path_fusion=args.path_fusion, self_fusion_weight=args.self_fusion_weight, ngpu=args.ngpu,
            mode=args.mode, apply_rescoring=args.apply_rescoring, is_visualized=args.is_visualized, output_dir=output_dir
        )
    elif model_type == 'fast-skip-transducer':
        return FastSkipTransducerRecognizer(
            model=model, lm=lm, lm_weight=args.lm_weight, beam_width=args.beam_width, max_len=args.max_len, nbest=args.nbest,
            idx2unit=idx2unit, path_fusion=args.path_fusion, self_fusion_weight=args.self_fusion_weight, one_step_constrain=args.one_step_constrain,
            ngpu=args.ngpu, mode=args.mode, blk_spike_threshold=args.blk_spike_threshold,
            is_expand_spikes=args.is_expand_spikes, left_expand=args.left_expand, right_expand=args.right_expand,
            is_visualized=args.is_visualized, output_dir=output_dir
        )    
    else:
        raise NotImplementedError
